import warnings
from csv import reader
from datetime import date, datetime
from urllib2 import urlopen
from copy import deepcopy
from pickle import dump
from time import strptime, sleep
from numpy import array, zeros, sum, dtype
from h5py import File
from BeautifulSoup import BeautifulSoup
from matplotlib.pyplot import figure, show, draw, ion
from pdb import set_trace

class StockError(Exception):
    """
    Exception to define errors in that occur
    in the Stock class
    """
    def __init__(self, message):
         self.message = message
    def __str__(self):
         return repr(self.message)

class Stock(object):
    """
    Base object for importing and storing stock information.
    """
    def __init__(self,ticker):
        self.ticker = ticker.lower().strip()            # string of ticker symbol
        self.quote = {
                        'price' : None,
                        'close' : None,
                        'open' : None,
                        'volume' : None,
                        'bid' : None,
                        'ask' : None,
                        'date' : None
                        }
        self.historical_data = zeros(0, dtype=[('date',date), ('open','f4'),
                        ('high','f4'), ('low','f4'), ('close','f4'),
                        ('volume','f4'), ('adj close','f4')])
        self.description = None         # dictionary of stock descriptors (e.g. full name, market segment, ...)
        self.fundamentals = None        # dictionary of fundamentals data (e.g. PE, alpha, beta, ...)
        
    def import_quote(self):
        """
        Get the market data (last close, open, current price, volume, ...)
        """
        quote = {}
        url = "http://finance.yahoo.com/q?s=%s"%(self.ticker)

        html = BeautifulSoup(urlopen(url))
        price = html.find("span",{"id":"yfs_l84_%s"%self.ticker}).getText()
        quote.update({'price' : float(price)})

        today = date.today()
        pdate = html.find("span",{"id":"yfs_t53_%s"%self.ticker}).getText().strip()
        if pdate.find('EDT') != -1:
            pdate = pdate.split('EDT')[0].strip()
        if pdate.find(',') != -1:
            pdate = strptime(pdate,"%b %d, %H:%M%p")
            pdate = datetime(today.year, pdate.tm_mon, pdate.tm_mday,
                            pdate.tm_hour, pdate.tm_min)
        elif (pdate.find('AM') != -1 or pdate.find('PM') != -1):
            pdate = strptime(pdate,"%H:%M%p")
            pdate = datetime(today.year, today.month, today.day,
                            pdate.tm_hour, pdate.tm_min)
        else:
            pdate = strptime(pdate,"%b %d")
            pdate = datetime(today.year, pdate.tm_mon, pdate.tm_mday)

        quote.update({'date' : pdate})

        try:
            table1 = html.find("table",{"id":"table1"})
            table2 = html.find("table",{"id":"table2"})
            for t in [table1,table2]:
                for tr in t.findAll("tr"):
                    key = str(tr.find("th").getText())
                    val = tr.find("td").getText()
                    if key.find('Open') != -1:
                        quote.update({'open' : float(val)})
                    elif key.find('Prev Close') != -1:
                        quote.update({'close' : float(val)})
                    elif key.find('Volume') != -1:
                        val = val.replace(',','')
                        quote.update({'volume' : float(val)})
                    elif key.find('Bid') != -1: 
                        val = val.split('x')[0]
                        quote.update({'bid' : float(val)})
                    elif key.find('Ask') != -1: 
                        val = val.split('x')[0]
                        quote.update({'ask' : float(val)})
        except:
            warnings.warn("Could not import all quote data for %s"%self.ticker)

        if set(quote.keys()).issubset(set(self.quote.keys())):
            self.quote.update(quote)
        else:
            raise StockError("Possible erroneous quote information imported for %s."
                    %(self.ticker))

    def import_historical_data(self, startdate, enddate):
        """
        Get historcal price and volume data for the analysis.  The data could
        already have been retrieved and stored within the variable, or a file.
        If data already exists, inspect the dates and ensure that the dates
        cover the requested range. If the requested data exists, do nothing.
        """
        if len(self.historical_data) == 0:
            self.historical_data = self.__import_historical_data(startdate, 
                                        enddate)
        elif (startdate < self.historical_data['date'][0] or 
              enddate > self.historical_data['date'][-1]):
            self.historical_data = self.__import_historical_data(startdate, 
                                        enddate)

    def __import_historical_data(self, startdate, enddate):
        """
        Get historcal price and volume data from yahoo finance.
        Furture code should have imports available from alternative
        web sources, and this function should simply be the access point
        for the other imports.
        """
        yahoo_month = ['00','01','02','03','04','05','06','07','08',
                       '09','10','11']
        url_base = "http://ichart.finance.yahoo.com/table.csv?"
        url = url_base+"s=%s&a=%s&b=%i&c=%i&d=%s&e=%i&f=%i&g=d&ignore=.csv" %(self.ticker,
                yahoo_month[startdate.month-1], startdate.day, startdate.year,
                yahoo_month[enddate.month-1], enddate.day, enddate.year)

        csv_reader = reader(urlopen(url))
        sleep(0.1) # pause while page loads
        header = csv_reader.next()
        header[0] = 'Date'
        quotes = [[date(*strptime(r[0],'%Y-%m-%d')[:3]),
             float(r[1]), float(r[2]), float(r[3]), float(r[4]), 
             float(r[5]), float(r[6])] for r in csv_reader]

        if quotes[-1][0] != startdate or quotes[0][0] != enddate:
            raise StockError("Valid date range for %s: %s to %s" 
                    %(self.ticker, quotes[-1][0], quotes[0][0]))

        historical_data = zeros(len(quotes), dtype=[('date',date), 
            ('open','f4'), ('high','f4'), ('low','f4'), ('close','f4'),
            ('volume','f4'), ('adj close','f4')])
        quotes = array(quotes)
        for i, nm in enumerate(['date','open','high','low','close',
                                'volume','adj close']):
            historical_data[nm] = quotes[::-1,i]
        return historical_data

    def get_quotes(self, startdate, enddate=None):
        """
        Get a subset of historical data already imported.
        If enddate is not given then a only a single
        day (startdate is returned).
        startdate - date object
        enddate - optional date object
        return numpy array of historical type
        """
        dates = self.historical_data['date']
        st_ind = dates.tolist().index(startdate) 
        if enddate != None:
            ed_ind = dates.tolist().index(enddate) + 1
        else:
            ed_ind = st_ind + 1
        pts = ed_ind - st_ind
        quotes = zeros(pts, dtype=self.historical_data.dtype)
        quotes = self.historical_data[st_ind:ed_ind]
        return quotes

    def subset_dates(self, startdate, enddate):
        """
        Return an equivalent stock object with
        a subset of the date range.  This is very
        useful for testing purposes.
        return stock object
        """
        stk = deepcopy(self)
        quotes = self.get_quotes(startdate, enddate)
        stk.historical_data = quotes
        return stk

    def import_fundamentals(self):
        """
        Get the fundamentals data
        """
        data = {}
        url = "http://finance.yahoo.com/q?s=%s"%(self.ticker)
        html = BeautifulSoup(urlopen(url))

        table1 = html.find("table",{"id":"table1"})
        table2 = html.find("table",{"id":"table2"})
        for t in [table1,table2]:
            for tr in t.findAll("tr"):
                key = str(tr.find("th").getText())
                val = tr.find("td").getText()
                if key.find('Market Cap') != -1:
                    if key.find('M'):
                        val = val.replace('M','')
                        val = float(val) * 1e6
                    if key.find('B'):
                        val = val.replace('B','')
                        val = float(val) * 1e9
                    data.update({'market cap' : val})
                elif key.find('P/E') != -1:
                    data.update({'P/E' : float(val)})
                elif key.find('EPS') != -1:
                    data.update({'EPS' : float(val)})
                elif key.find('Yield') != -1:
                    val = val.split('(')[0]
                    data.update({'divident' : float(val)})
                elif key.find('Beta') != -1:
                    data.update({'beta' : float(val)})

        self.fundamentals = data

    def import_description(self):
        """
        Lookup details of the stock
        from the internet
        # TODO
        """
        pass

    def save(self, fname, hdf5_path):
        """
        Open the hdf5 file with string name fname and
        write the object data.
        fname - string filename of HDF5 file 
        hdf5_path - HDF5 file path where new group will be created
        """
        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'

        name = self.ticker + '/'

        full_name = hdf5_path + name
        
        with File(fname, mode='a') as f_h5py:
            f_h5py.create_group(full_name)
            
            hist_path = full_name + 'historical_data/'
            f_h5py.create_group(hist_path)
            for nm in self.historical_data.dtype.names:
                if nm == 'date':
                    f_h5py.create_dataset(hist_path+nm, 
                            data=map(date.toordinal, self.historical_data[nm]))
                else:
                    f_h5py.create_dataset(hist_path+nm, 
                            data=self.historical_data[nm])

            quote_path = full_name + 'quotes/'
            f_h5py.create_group(quote_path)
            for k, v in self.quote.iteritems():
                if k == 'date':
                    f_h5py.create_dataset(quote_path+k, data=v.toordinal())
                elif v is None:
                    f_h5py.create_dataset(quote_path+k, data='N/A')
                else:
                    f_h5py.create_dataset(quote_path+k, data=v)
 
    def load(self, fname, hdf5_path):
        """
        Open the hdf5 file with string name fname
        and read the data from the hdf5 file.
        fname - string file name of h5py file
        hdf5_path - root path that contains the stock
                    data
        """
        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'

        self.ticker = hdf5_path.split('/')[-2]

        with File(fname, mode='r') as f_h5py:
            hist_path = hdf5_path + 'historical_data/'
            length = f_h5py[hist_path+'date'].shape[0]
            historical_data = zeros((length), dtype=self.historical_data.dtype)
            for nm in self.historical_data.dtype.names:
                if nm == 'date':
                    historical_data[nm] = array(map(date.fromordinal,
                           f_h5py[hist_path+nm].value))
                else:
                    historical_data[nm] = array(f_h5py[hist_path+nm].value)
            self.historical_data = historical_data

            quote_path = hdf5_path + 'quotes/'
            for k in self.quote.iterkeys():
                if k == 'date':
                    self.quote.update({k : date.fromordinal(f_h5py[quote_path+k].value)})
                elif f_h5py[quote_path+k].value == 'N/A':
                    self.quote.update({k : None})
                else:
                    self.quote.update({k : f_h5py[quote_path+k].value})

    def plot(self):
        """
        Plot the historic price data
        """
        ion()
        fig = figure()
        ax = fig.add_subplot(111)
        for p, c in [['close','blue'], ['open','red'], ['high','green'], 
                ['low','brown'], ['adj close','gold']]:
            ax.plot(self.historical_data['date'], self.historical_data[p], 
                color=c, label=p)
        ax.set_xlabel('Date')
        ax.set_ylabel('price')
        ax.set_title(self.ticker)
        ax.legend()
        draw()
        show()

class StockCollection(dict):
    """
    Collection of stocks for constructing portfolios,
    markets, ...etc.
    """
    def __init__(self, tickers=None):
        """
        Initialize the stock collection dictionary with
        a passed list of ticker symbols
        tickers - list of ticker symbol strings
        """
        if tickers is not None:
            for t in tickers:
                t = t.lower().strip()
                self.update({t : Stock(t)})

    @classmethod
    def from_array(cls, stkarray):
        """
        Construct a StockCollection from an array
        """
        coll = cls()

        for s in stkarray:
            coll.add(s['ticker'])

            for k in ['ask', 'bid', 'close', 'date', 'open', 'price', 'volume']:
                coll[s['ticker']].quote[k] = s['quote'][k]

            coll[s['ticker']].historical_data = s['historical data']

        return coll

    def add(self, ticker):
        """
        Add a stock to the collection
        """
        ticker = ticker.lower().strip()
        self.update({ticker : Stock(ticker)})

    def remove(self, ticker):
        """
        Remove a stock from the collection
        """
        del self[ticker]

    def import_quotes(self):
        """
        Update the price quote for each stock in the collection.
        """
        errors = []
        for k, v in self.iteritems():
            try:
                v.import_quote()
            except Exception as e:
                mesg = """Could not update quote for %s.
                    %s""" %(k, e.args)
                warnings.warn(mesg)
                errors.append([k, e.args])
    
        if len(errors) != 0:
            raise StockError("""Errors updating quotes for %s."""%(errors))

    def import_historical_data(self, startdate, enddate):
        """
        Get historical data for each stock in the collection.
        """
        errors = []
        for k, v in self.iteritems():
            try:
                v.import_historical_data(startdate, enddate)
            except Exception as e:
                mesg = "Could not import historical data for %s.\n%s" %(k, e.args)
                warnings.warn(mesg)
                errors.append(k)

        if len(errors) != 0:
            raise StockError("Errors importing historical data for %s. \n%s"%(k, e.args))

    def import_fundamentals(self):
        """
        Get the fundamentals for each stock in the collection
        """
        errors = []
        for k, v in self.iteritems():
            try:
                v.import_fundamentals()
            except Exception as e:
                mesg = "Could not get fundamentals data for %s.\n%s" %(k, e.args)
                warnings.warn(mesg)
                errors.append(k)

        if len(errors) != 0:
            raise StockError("Errors updating quotes for %s. \n%s"%(k, e.args))

    def subset_dates(self, startdate, enddate):
        """
        Return an equivalent stock collection object with
        a subset of the date range.  This is very
        useful for testing purposes.
        return stockCollrection object
        """
        col = deepcopy(self)
        for stk in col.itervalues():
            quotes = stk.get_quotes(startdate, enddate)
            stk.historical_data = quotes
        return col

    def to_array(self):
        """
        Make an array for each historical_data quantity
        containing data from all the stocks in the
        collection. (This is to facilitate some calculations)
        """
        # get the number of date points for one of the holdings
        n = self.get(self.keys()[0]).historical_data['close'].shape[0]
        m = len(self)
        quote_dtype = dtype([('ask','<f4'), ('bid','<f4'),
            ('close','<f4'), ('date',date), ('open','<f4'), ('price','<f4'), 
            ('volume','<f4')])
        hist_dtype = dtype([('date',date, n), ('open','f4',n), 
            ('high','f4',n), ('low','f4',n), ('close','f4',n), 
            ('volume','f4',n), ('adj close','f4',n)])
        price_array = zeros((m), dtype=[('ticker','|S5'), ('quote',quote_dtype),
                        ('historical data',hist_dtype)])

        for i, s in enumerate(self):
            price_array['ticker'][i] = s

            for k in ['date', 'ask', 'bid', 'close', 'open', 'price', 'volume']:
                price_array['quote'][k] = self[s].quote[k]

            for k in ['date', 'open','high','low','close','volume','adj close']:
                price_array['historical data'][k][i] = self[s].historical_data[k]
        return price_array
        
    def save(self, fname, hdf5_path):
        """
        Save the stock collection to an HDF5 file
        with all stocks saved in subgrouops of the
        hdf5_path.
        fname - string filename of hdf5 file
        hdf5_path - root path of hdf5 file to create
                groups for each stock
        """
        for k, v in self.iteritems():
            v.save(fname, hdf5_path)

    def load(self, fname, hdf5_path):
        """
        Load the data for the stock collection from
        an HDF5 file with each stock contained in a
        subtroup of hdf5_path.
        fname - string filename of hdf5 file
        hdf5_path - root path of hdf5 file where
                each stock subgroup is stored.
        """
        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'

        with File(fname, mode='r') as f_h5py:
            tickers = f_h5py[hdf5_path].keys()
        tickers = map(str,tickers)

        for ticker in tickers:
            self.update({ticker : Stock(ticker)})
            full_path = hdf5_path + ticker
            self[ticker].load(fname, full_path)

if __name__ == '__main__':
    """
    Create a StockCollection from the command line arguments
    Test the stock.get_quote(date) method
    Plot the stock.plot_close() method
    example:
    $python stock.py mmm ge pfe
    """
    import sys
    startdate = date(2012,7,23)
    enddate = date(2013,7,22)
    coll = StockCollection()
    for ticker in sys.argv[1:]:
        coll.add(ticker)
    coll.import_quotes()
    coll.import_historical_data(startdate,enddate)
    coll.import_fundamentals()
    quote_array = coll.make_array()
    stk = coll[sys.argv[1]]
    quote_date = date(2013,2,6)
    close = stk.get_quotes(quote_date)['close']
    stk.plot()
